#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :  db_tools.py
@Time    :  2023/09/28
@Author  :  HDUZN
@Version :  1.0
@Contact :  hduzn@vip.qq.com
@License :  (C)Copyright 2023-2024
@Desc    :  创建数据库
'''

import sqlite3, os
import pandas as pd
import numpy as np

# 创建数据库
def create_table(db_file):    
    # 如果文件存在，就删除它
    if os.path.exists(db_file):
        os.remove(db_file)

    # 创建数据库文件
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    
    # 创建车次信息表
    sql = '''
        CREATE TABLE "trains" (
        "id"    INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT UNIQUE,
        "start_city"  TEXT NOT NULL,
        "end_city"  TEXT NOT NULL,
        "train_no" TEXT NOT NULL,
        "train_info" TEXT NOT NULL,
        "time_info" TEXT NOT NULL,
        "seat_first"  REAL NOT NULL,
        "seat_second"  REAL NOT NULL
        )
    '''
    cursor.execute(sql)

    cursor.close()
    conn.commit()
    conn.close()

# 导入Excel表格数据到SQLite数据库
def insert_scores_to_db(data_file, db_file):
    # 创建SQLite数据库连接
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()

    # 读取Excel表中的数据
    df = pd.read_excel(data_file, sheet_name='Sheet1')
    data_list = np.array(df).tolist()
    # print(data_list)

    # 插入数据库
    sql = 'INSERT INTO trains (start_city,end_city,train_no,train_info,time_info,seat_first,seat_second) VALUES (?,?,?,?,?,?,?)'
    cursor.executemany(sql, data_list)

    # 提交事务并关闭连接
    conn.commit()
    conn.close()

def main():
    # 创建数据库
    db_file = r'.\static\db\ticket.db'
    data_file = r'.\static\ex_data\data.xlsx'

    # 创建数据库
    create_table(db_file)
    # 数据导入
    insert_scores_to_db(data_file, db_file)

    print('-'*10, 'fine!')

main()
